from typing import Generic, Iterable, Mapping, TypeVar, Union

from zkl_pyutils_registries.basics import Matcher, SpecialKey, UNSET

AnyTag = TypeVar("AnyTag")
AnyKey = TypeVar("AnyKey")
AnyQuery = TypeVar("AnyQuery")
AnyMatch = TypeVar("AnyMatch")

ComponentMatcher = Matcher[Union[AnyKey, SpecialKey], Union[AnyKey, SpecialKey], AnyMatch]
ComposedQuery = Union[Mapping[AnyTag, AnyKey], Iterable[tuple[AnyTag, AnyKey]]]
ComposedKey = Union[Mapping[AnyTag, AnyKey], Iterable[tuple[AnyTag, AnyKey]]]
ComposedMatch = tuple[AnyMatch, ...]


class ComposedMatcher(
    dict[AnyTag, ComponentMatcher],
    Matcher[ComposedQuery, ComposedKey, ComposedMatch],
    Generic[AnyTag, AnyQuery, AnyKey, AnyMatch],
):
    def match(self, query: ComposedQuery, key: ComposedKey) -> ComposedMatch:
        query = dict(query)
        key = dict(key)
        return tuple(
            matcher.match(query.get(tag, UNSET), key.get(tag, UNSET))
            for tag, matcher in self.items())
